# -*- coding: utf-8 -*-
"""
@author:XuMing(xuming624@qq.com)
@description: Image ocr demo.
"""
import time
import sys

sys.path.append('..')
from imgocr import ImgOcr
from imgocr import draw_ocr_boxes
from PIL import Image

from deepdoc.vision import LayoutRecognizer

if __name__ == "__main__":
    m = ImgOcr(use_gpu=False, is_efficiency_mode=True)
    img_path = sys.argv[1]
    s = time.time()
    result = m.ocr(img_path)
    e = time.time()
    print("total time: {:.4f} s".format(e - s))
    print("result:", result)
    for i in result:
        print(i['text'])

    print("layout information")
    img = Image.open(img_path)
    detr = LayoutRecognizer("layout")
    layouts = detr.forward([img], thr=0.5)
    for i, lyt in enumerate(layouts):
        if lyt[0]['type'] != 'text':    # 显示非正文的内容
            print(lyt[0])
    # draw boxes
    # draw_ocr_boxes(img_path, result, '11_box.jpg')
    # print('Save result to 11_box.jpg')
